"""
使用训练好的模型进行对弈
支持人机对战和AI自我对弈
"""

import pygame
import sys
import os
from board import Board
from mcts_alpha import AlphaMCTS
from neural_network import GomokuNet
import torch


class TrainedGameGUI:
    """训练模型对弈界面"""
    
    def __init__(self, model_path='models/best_model.pth', num_simulations=100):
        """初始化"""
        # 加载模型
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model = GomokuNet(board_size=15).to(self.device)
        
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=self.device)
            self.model.load_state_dict(checkpoint['model_state_dict'])
            print(f"✅ 模型加载成功: {model_path}")
        else:
            print(f"⚠️  模型文件不存在: {model_path}")
            print("使用未训练的模型（随机策略）")
        
        self.model.eval()
        self.num_simulations = num_simulations
        
        # 初始化游戏
        self.board = Board(size=15)
        self.cell_size = 40
        self.margin = 50
        self.window_size = self.cell_size * 14 + self.margin * 2
        
        # 初始化Pygame
        pygame.init()
        self.screen = pygame.display.set_mode((self.window_size, self.window_size + 100))
        pygame.display.set_caption("五子棋 - AlphaGo Zero AI")
        self.clock = pygame.time.Clock()
        
        # 颜色
        self.COLOR_BG = (220, 179, 92)
        self.COLOR_LINE = (0, 0, 0)
        self.COLOR_BLACK = (0, 0, 0)
        self.COLOR_WHITE = (255, 255, 255)
        self.COLOR_TEXT = (0, 0, 0)
        self.COLOR_BUTTON = (100, 150, 200)
        self.COLOR_BUTTON_HOVER = (120, 170, 220)
        
        # 游戏状态
        self.game_over = False
        self.winner = None
        self.ai_thinking = False
        self.last_move = None
        
        # 按钮
        self.restart_button = pygame.Rect(self.margin, self.window_size + 20, 150, 40)
        self.mode_button = pygame.Rect(self.margin + 170, self.window_size + 20, 150, 40)
        
        # 游戏模式
        self.ai_vs_ai = False
    
    def draw_board(self):
        """绘制棋盘"""
        self.screen.fill(self.COLOR_BG)
        
        # 绘制网格
        for i in range(15):
            # 横线
            start_pos = (self.margin, self.margin + i * self.cell_size)
            end_pos = (self.margin + 14 * self.cell_size, self.margin + i * self.cell_size)
            pygame.draw.line(self.screen, self.COLOR_LINE, start_pos, end_pos, 2)
            
            # 竖线
            start_pos = (self.margin + i * self.cell_size, self.margin)
            end_pos = (self.margin + i * self.cell_size, self.margin + 14 * self.cell_size)
            pygame.draw.line(self.screen, self.COLOR_LINE, start_pos, end_pos, 2)
        
        # 绘制天元和星位
        star_positions = [(3, 3), (3, 11), (11, 3), (11, 11), (7, 7)]
        for row, col in star_positions:
            x = self.margin + col * self.cell_size
            y = self.margin + row * self.cell_size
            pygame.draw.circle(self.screen, self.COLOR_LINE, (x, y), 5)
    
    def draw_stones(self):
        """绘制棋子"""
        for row in range(15):
            for col in range(15):
                if self.board.board[row, col] != 0:
                    x = self.margin + col * self.cell_size
                    y = self.margin + row * self.cell_size
                    color = self.COLOR_BLACK if self.board.board[row, col] == 1 else self.COLOR_WHITE
                    pygame.draw.circle(self.screen, color, (x, y), self.cell_size // 2 - 2)
                    
                    # 最后一步标记
                    if self.last_move and (row, col) == self.last_move:
                        marker_color = self.COLOR_WHITE if color == self.COLOR_BLACK else self.COLOR_BLACK
                        pygame.draw.circle(self.screen, marker_color, (x, y), 5)
    
    def draw_ui(self):
        """绘制UI元素"""
        # 状态文本
        if self.game_over:
            if self.winner == 1:
                text = "黑棋获胜！"
            elif self.winner == -1:
                text = "白棋获胜！"
            else:
                text = "平局！"
        elif self.ai_thinking:
            text = "AI思考中..."
        else:
            player = "黑棋" if self.board.current_player == 1 else "白棋"
            text = f"当前: {player}"
        
        font = pygame.font.Font(None, 36)
        text_surface = font.render(text, True, self.COLOR_TEXT)
        text_rect = text_surface.get_rect(center=(self.window_size // 2, 20))
        self.screen.blit(text_surface, text_rect)
        
        # 重新开始按钮
        mouse_pos = pygame.mouse.get_pos()
        restart_color = self.COLOR_BUTTON_HOVER if self.restart_button.collidepoint(mouse_pos) else self.COLOR_BUTTON
        pygame.draw.rect(self.screen, restart_color, self.restart_button, border_radius=5)
        
        font_small = pygame.font.Font(None, 28)
        restart_text = font_small.render("重新开始", True, self.COLOR_WHITE)
        restart_rect = restart_text.get_rect(center=self.restart_button.center)
        self.screen.blit(restart_text, restart_rect)
        
        # 模式切换按钮
        mode_color = self.COLOR_BUTTON_HOVER if self.mode_button.collidepoint(mouse_pos) else self.COLOR_BUTTON
        pygame.draw.rect(self.screen, mode_color, self.mode_button, border_radius=5)
        
        mode_text_str = "AI vs AI" if not self.ai_vs_ai else "人机对战"
        mode_text = font_small.render(mode_text_str, True, self.COLOR_WHITE)
        mode_rect = mode_text.get_rect(center=self.mode_button.center)
        self.screen.blit(mode_text, mode_rect)
        
        # 显示模拟次数
        info_text = f"MCTS模拟: {self.num_simulations}次"
        info_surface = font_small.render(info_text, True, self.COLOR_TEXT)
        self.screen.blit(info_surface, (self.margin + 340, self.window_size + 30))
    
    def get_board_pos(self, mouse_pos):
        """将鼠标位置转换为棋盘坐标"""
        x, y = mouse_pos
        col = round((x - self.margin) / self.cell_size)
        row = round((y - self.margin) / self.cell_size)
        
        if 0 <= row < 15 and 0 <= col < 15:
            return row, col
        return None
    
    def ai_move(self):
        """AI落子"""
        self.ai_thinking = True
        
        # 使用AlphaMCTS
        mcts = AlphaMCTS(
            board=self.board,
            neural_net=self.model,
            num_simulations=self.num_simulations,
            c_puct=1.5
        )
        
        move = mcts.get_best_move(self.board.current_player, temperature=0)
        
        if move:
            self.board.place_stone(move[0], move[1])
            self.last_move = move
            
            # 检查胜负
            winner = self.board.check_winner()
            if winner is not None:
                self.game_over = True
                self.winner = winner
        
        self.ai_thinking = False
    
    def restart_game(self):
        """重新开始游戏"""
        self.board = Board(size=15)
        self.game_over = False
        self.winner = None
        self.last_move = None
        self.ai_thinking = False
    
    def run(self):
        """主循环"""
        running = True
        
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                
                elif event.type == pygame.MOUSEBUTTONDOWN and not self.ai_thinking:
                    mouse_pos = pygame.mouse.get_pos()
                    
                    # 检查按钮点击
                    if self.restart_button.collidepoint(mouse_pos):
                        self.restart_game()
                    elif self.mode_button.collidepoint(mouse_pos):
                        self.ai_vs_ai = not self.ai_vs_ai
                        self.restart_game()
                    
                    # 人机对战模式下的玩家落子
                    elif not self.game_over and not self.ai_vs_ai:
                        pos = self.get_board_pos(mouse_pos)
                        if pos and self.board.is_valid_move(pos[0], pos[1]):
                            self.board.place_stone(pos[0], pos[1])
                            self.last_move = pos
                            
                            # 检查胜负
                            winner = self.board.check_winner()
                            if winner is not None:
                                self.game_over = True
                                self.winner = winner
            
            # AI vs AI 模式
            if self.ai_vs_ai and not self.game_over and not self.ai_thinking:
                self.ai_move()
                pygame.time.wait(500)  # 延迟以便观察
            
            # 人机对战模式 - AI自动落子
            elif not self.ai_vs_ai and not self.game_over and not self.ai_thinking:
                if self.board.current_player == -1:  # 白棋是AI
                    self.ai_move()
            
            # 绘制
            self.draw_board()
            self.draw_stones()
            self.draw_ui()
            
            pygame.display.flip()
            self.clock.tick(60)
        
        pygame.quit()
        sys.exit()


if __name__ == "__main__":
    print("=" * 60)
    print("🎮 五子棋 - AlphaGo Zero AI 对弈")
    print("=" * 60)
    print()
    
    # 检查模型
    model_path = 'models/best_model.pth'
    if os.path.exists(model_path):
        print(f"✅ 找到训练模型: {model_path}")
    else:
        print(f"⚠️  未找到训练模型: {model_path}")
        print("将使用未训练的模型（随机策略）")
        print()
        print("提示: 运行 'python train.py' 进行训练")
    
    print()
    print("请选择MCTS模拟次数:")
    print("1. 快速 (50次) - 约0.5秒/步")
    print("2. 标准 (100次) - 约1秒/步")
    print("3. 强力 (200次) - 约2秒/步")
    print("4. 极限 (500次) - 约5秒/步")
    print()
    
    choice = input("请输入选择 (1-4, 默认2): ").strip()
    
    simulations_map = {'1': 50, '2': 100, '3': 200, '4': 500}
    num_simulations = simulations_map.get(choice, 100)
    
    print()
    print(f"🚀 启动游戏 (MCTS模拟: {num_simulations}次)")
    print()
    print("游戏说明:")
    print("- 点击棋盘落子")
    print("- 点击'重新开始'重置游戏")
    print("- 点击'AI vs AI'切换模式")
    print()
    
    # 启动游戏
    game = TrainedGameGUI(model_path=model_path, num_simulations=num_simulations)
    game.run()
